# 导入必要的库
import gradio as gr
import numpy as np
import joblib
import cv2
import logging
from pathlib import Path
import glob

# 使用独立的 Keras 3.x
from keras.applications.vgg16 import VGG16, preprocess_input
from keras.utils import img_to_array

# 配置logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 加载 VGG16 模型用于特征提取
logging.info("正在加载 VGG16 模型...")
vgg_model = VGG16(weights='imagenet', include_top=False, pooling="max")
logging.info("VGG16 模型加载完成")

# 加载训练好的最佳分类模型
def load_best_model():
    """加载保存的最佳模型"""
    # 查找所有保存的模型文件
    model_files = glob.glob('best_model_*.pkl')
    
    if not model_files:
        raise FileNotFoundError("未找到保存的模型文件。请先运行 train.py 训练模型。")
    
    # 使用第一个找到的模型文件
    model_path = model_files[0]
    logging.info(f"正在加载模型: {model_path}")
    
    model = joblib.load(model_path)
    
    # 尝试加载模型信息
    try:
        model_info = joblib.load('model_info.pkl')
        logging.info(f"模型名称: {model_info['model_name']}")
        logging.info(f"模型准确率: {model_info['accuracy']:.4f}")
        return model, model_info
    except:
        logging.warning("未找到模型信息文件")
        return model, None

# 加载模型
try:
    best_model, model_info = load_best_model()
    logging.info("模型加载成功！")
except Exception as e:
    logging.error(f"加载模型失败: {e}")
    logging.error("请先运行 train.py 训练模型")
    best_model = None
    model_info = None

def extract_features(image):
    """使用 VGG16 提取图像特征"""
    # 调整图像大小到 224x224
    img = cv2.resize(image, (224, 224))
    
    # 转换为 RGB（如果是灰度图）
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    elif img.shape[2] == 4:
        img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
    else:
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    # 转换为数组并预处理
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    
    # 提取特征
    features = vgg_model.predict(img_array, verbose=0)
    
    return features.flatten()

def predict_image(image):
    """预测图像是猫还是狗"""
    if best_model is None:
        return {
            "错误": 1.0
        }, "❌ 模型未加载，请先运行 train.py 训练模型"
    
    try:
        # 提取特征
        features = extract_features(image)
        features = features.reshape(1, -1).astype('float32')
        
        # 进行预测
        prediction = best_model.predict(features)[0]
        
        # 获取预测概率（如果模型支持）
        if hasattr(best_model, 'predict_proba'):
            proba = best_model.predict_proba(features)[0]
            cat_prob = proba[0]
            dog_prob = proba[1]
        else:
            # 如果模型不支持概率预测，使用硬分类
            cat_prob = 1.0 if prediction == 0 else 0.0
            dog_prob = 1.0 if prediction == 1 else 0.0
        
        # 确定预测结果
        result = "🐶 狗 (Dog)" if prediction == 1 else "🐱 猫 (Cat)"
        
        # 构建置信度字典
        confidence = {
            "🐱 猫 (Cat)": float(cat_prob),
            "🐶 狗 (Dog)": float(dog_prob)
        }
        
        # 构建详细信息
        if model_info:
            info = f"""
### 预测结果: {result}

**模型信息:**
- 模型名称: {model_info.get('model_name', 'Unknown')}
- 模型准确率: {model_info.get('accuracy', 0):.2%}
- 测试集准确率: {model_info.get('test_accuracy', 0):.2%}

**预测置信度:**
- 猫的概率: {cat_prob:.2%}
- 狗的概率: {dog_prob:.2%}
            """
        else:
            info = f"""
### 预测结果: {result}

**预测置信度:**
- 猫的概率: {cat_prob:.2%}
- 狗的概率: {dog_prob:.2%}
            """
        
        return confidence, info
        
    except Exception as e:
        logging.error(f"预测失败: {e}")
        return {
            "错误": 1.0
        }, f"❌ 预测失败: {str(e)}"

# 创建 Gradio 界面
def create_interface():
    """创建 Gradio Web 界面"""
    
    # 创建示例图片列表
    examples = []
    test_folder = Path("data/test")
    if test_folder.exists():
        # 获取前5张测试图片作为示例
        test_images = list(test_folder.glob("*.jpg"))[:5]
        examples = [[str(img)] for img in test_images]
    
    # 界面标题和描述
    title = "🐱🐶 猫狗分类器 - LazyPredict 版本"
    
    description = """
    ## 使用 LazyPredict 训练的最佳模型进行猫狗分类
    
    上传一张猫或狗的图片，AI 将告诉你这是猫还是狗！
    
    **使用说明:**
    1. 点击上传图片或拖拽图片到上传区域
    2. 或者点击下方的示例图片进行快速测试
    3. 查看预测结果和置信度
    """
    
    if model_info:
        description += f"""
    
    **当前使用的模型:** {model_info.get('model_name', 'Unknown')}
    
    **模型性能:** 准确率 {model_info.get('accuracy', 0):.2%}
        """
    
    # 创建界面
    interface = gr.Interface(
        fn=predict_image,
        inputs=gr.Image(label="上传图片", type="numpy"),
        outputs=[
            gr.Label(label="预测置信度", num_top_classes=2),
            gr.Markdown(label="详细信息")
        ],
        title=title,
        description=description,
        examples=examples,
        theme=gr.themes.Soft(),
        allow_flagging="never"
    )
    
    return interface

# 主函数
def main():
    if best_model is None:
        logging.error("模型未加载，无法启动 Web 应用")
        logging.error("请先运行: python train.py")
        return
    
    # 创建并启动界面
    interface = create_interface()
    logging.info("正在启动 Gradio Web 应用...")
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )

if __name__ == "__main__":
    main()
